# -*- coding: utf-8 -*-
'''
File:modeling_barranzi2VL.py                                         
Time:2025/9/29 15:43                                    
IDE:PyCharm                                     
Author:Barranzi An                                        
email:awc19930818@outlook.com                           
github:https://github.com/La0bALanG                     
Barranzi's Blog:私聊巴郎子要文档链接                  
requirement:(Please describle your requirement here) -->
'''

from typing import Optional
import torch
from torch import nn

from .modeling_barranzi2LM import RMSNorm, Barranzi2LMForCausalLM
from .configuration_barranzi2 import VLMConfig


class MultiModalProjector(nn.Module):
    """
    多模态投影器类，用于将视觉特征投影到语言模型的特征空间中。

    该类实现了视觉-语言模型中的关键模态对齐功能，通过空间池化和线性投影，
    将高分辨率的视觉特征转换为与语言模型兼容的token序列表示。主要应用于
    多模态大模型中，实现图像特征与文本特征的统一表示。

    核心功能：
    1. 空间下采样：通过平均池化将密集的图像块特征降采样到目标token数量
    2. 特征归一化：使用RMSNorm对视觉特征进行标准化处理
    3. 维度投影：通过线性变换将视觉特征维度映射到语言模型隐藏层维度

    技术特点：
    - 支持动态配置的图像尺寸、块大小和目标token数量
    - 采用平均池化进行空间聚合，保留全局语义信息
    - 使用可学习的投影权重实现模态间的特征对齐
    - 保持数值精度的一致性，支持不同dtype的输入输出

    应用场景：
    - 视觉-语言预训练模型的模态对齐模块
    - 多模态理解任务中的特征融合
    - 图像描述生成、视觉问答等下游任务
    """

    def __init__(self, config: VLMConfig):
        """
        初始化多模态投影器。

        参数：
            config (VLMConfig): 视觉-语言模型配置对象，包含以下关键配置：
                - vision_hidden_size: 视觉编码器的隐藏层维度
                - hidden_size: 语言模型的隐藏层维度
                - image_size: 输入图像的尺寸（正方形边长）
                - patch_size: 视觉编码器的图像块大小
                - tokens_per_image: 每张图像目标输出的token数量

        初始化过程：
        1. 创建投影权重矩阵，用于视觉特征到语言特征空间的线性变换
        2. 初始化视觉特征归一化层
        3. 计算空间池化相关参数（每图像块数、每边token数、池化核大小）
        4. 创建平均池化层用于空间下采样
        """
        super().__init__()

        # 投影权重矩阵：将视觉特征维度映射到语言模型隐藏层维度
        # 形状：(vision_hidden_size, hidden_size)
        self.input_projection_weight = nn.Parameter(
            torch.zeros(config.vision_hidden_size, config.hidden_size)
        )

        # 视觉特征归一化层：使用RMSNorm对投影前的视觉特征进行标准化
        self.vision_norm = RMSNorm(config.vision_hidden_size)

        # 计算空间下采样相关参数
        self.patches_per_image = int(config.image_size // config.patch_size)  # 每边图像块数量
        self.tokens_per_side = int(config.tokens_per_image ** 0.5)  # 每边目标token数量
        self.kernel_size = self.patches_per_image // self.tokens_per_side  # 池化核大小

        # 平均池化层：用于将密集的图像块特征下采样到目标token数量
        # kernel_size和stride相同，确保无重叠的池化操作
        self.avg_pool = nn.AvgPool2d(kernel_size=self.kernel_size, stride=self.kernel_size)

    def forward(self, vision_outputs: torch.Tensor):
        """
        前向传播：将视觉编码器输出投影为语言模型可用的特征序列。

        参数：
            vision_outputs (torch.Tensor): 视觉编码器输出的特征张量
                形状：(batch_size, patches_per_image*patches_per_image, vision_hidden_size)
                其中 patches_per_image = image_size // patch_size

        返回：
            torch.Tensor: 投影后的视觉特征，可直接用于语言模型
                形状：(batch_size, tokens_per_image, hidden_size)
                其中 tokens_per_image = tokens_per_side * tokens_per_side

        处理流程：
        1. 维度重排：将序列维度转换为空间维度，为2D池化做准备
        2. 空间池化：通过平均池化降低空间分辨率，减少token数量
        3. 维度恢复：将空间维度转换回序列维度
        4. 特征归一化：对池化后的特征进行标准化处理
        5. 线性投影：将视觉特征投影到语言模型的特征空间
        6. 类型对齐：确保输出类型与输入类型一致

        技术细节：
        - 使用transpose和reshape实现维度变换，保持内存连续性
        - 平均池化操作保留局部区域的语义信息，实现空间聚合
        - 投影权重支持动态dtype转换，确保数值精度匹配
        - 最终输出类型与输入类型对齐，保持计算图的一致性
        """
        # 输入张量形状：(batch_size, patches_per_image*patches_per_image, vision_hidden_size)
        batch_size, _, vision_hidden_size = vision_outputs.shape

        # 步骤1：维度重排 - 将序列维度转换为通道维度，为2D池化做准备
        # (batch_size, vision_hidden_size, patches_per_image*patches_per_image)
        reshaped_vision_outputs = vision_outputs.transpose(1, 2)
        # (batch_size, vision_hidden_size, patches_per_image, patches_per_image)
        reshaped_vision_outputs = reshaped_vision_outputs.reshape(
            batch_size, vision_hidden_size, self.patches_per_image, self.patches_per_image
        )
        reshaped_vision_outputs = reshaped_vision_outputs.contiguous()  # 确保内存连续性

        # 步骤2：空间池化 - 通过平均池化降低空间分辨率
        # (batch_size, vision_hidden_size, tokens_per_side, tokens_per_side)
        pooled_vision_outputs = self.avg_pool(reshaped_vision_outputs)
        # (batch_size, vision_hidden_size, tokens_per_side*tokens_per_side)
        pooled_vision_outputs = pooled_vision_outputs.flatten(2)
        # (batch_size, tokens_per_side*tokens_per_side, vision_hidden_size)
        pooled_vision_outputs = pooled_vision_outputs.transpose(1, 2)

        # 步骤3：特征归一化 - 对池化后的视觉特征进行标准化
        normed_vision_outputs = self.vision_norm(pooled_vision_outputs)

        # 步骤4：线性投影 - 将视觉特征映射到语言模型特征空间
        # (batch_size, tokens_per_side*tokens_per_side, hidden_size)
        projected_vision_outputs = torch.matmul(
            normed_vision_outputs,
            self.input_projection_weight.to(normed_vision_outputs.dtype)  # 确保数值精度匹配
        )

        # 步骤5：类型对齐 - 确保输出类型与原始输入类型一致
        return projected_vision_outputs.type_as(vision_outputs)


class Barranzi2VL(Barranzi2LMForCausalLM):
    """
    视觉-语言模型类，继承自LlmModel，实现多模态理解功能。

    该类是视觉-语言多模态模型的核心实现，通过集成视觉编码器和多模态投影器，
    将图像特征与文本特征统一到同一语义空间中，实现跨模态的理解和生成能力。
    作为LlmModel的扩展，它保留了语言模型的全部功能，同时增加了视觉处理能力。

    核心架构：
    1. 视觉编码器（Vision Tower）：负责提取图像的视觉特征表示
    2. 多模态投影器（MultiModalProjector）：将视觉特征投影到语言模型特征空间
    3. 语言模型主干（LlmModel）：处理文本序列并实现跨模态融合

    技术特点：
    - 模态对齐：通过可学习的投影权重实现视觉-语言特征空间的语义对齐
    - 灵活配置：支持不同的视觉编码器架构和投影策略
    - 统一表示：图像和文本在统一的特征空间中进行交互
    - 兼容性：完全兼容纯文本语言模型的接口和功能

    应用场景：
    - 图像描述生成（Image Captioning）
    - 视觉问答（Visual Question Answering）
    - 多模态对话（Multimodal Conversation）
    - 图文理解与推理（Vision-Language Understanding）
    - 跨模态检索（Cross-modal Retrieval）

    设计原理：
    采用"编码器-投影器-解码器"的架构范式，其中视觉编码器提取图像特征，
    投影器实现模态对齐，语言模型作为统一的语义空间处理跨模态信息。
    这种设计既保持了各模态处理的专业性，又实现了特征层面的深度融合。
    """

    def __init__(self, config: VLMConfig):
        """
        初始化视觉-语言模型。

        参数：
            config (VLMConfig): 视觉-语言模型配置对象，包含：
                - vision_tower: 视觉编码器实例（如ViT、CLIP等）
                - vision_hidden_size: 视觉特征维度
                - hidden_size: 语言模型隐藏层维度
                - image_size: 输入图像尺寸
                - patch_size: 视觉编码器的图像块大小
                - tokens_per_image: 每张图像的目标token数量
                - image_tok: 图像特殊token的ID
                - padding_tok: 填充token的ID
                - 以及继承自LlmModel的所有配置参数

        初始化流程：
        1. 调用父类LlmModel的初始化，建立语言模型基础架构
        2. 保存VLM配置，用于多模态相关操作
        3. 设置视觉编码器，负责图像特征提取
        4. 初始化多模态投影器，实现视觉-语言特征对齐
        """
        super().__init__(config)  # 初始化父类LlmModel，建立语言模型架构
        self.vlm_config = config  # 保存VLM专用配置
        self.vision_tower = config.vision_tower  # 视觉编码器（如ViT、CLIP等）
        self.multi_modal_projector = MultiModalProjector(config)  # 多模态投影器

    def get_image_features(self, pixel_values: torch.Tensor) -> torch.Tensor:
        """
        提取图像特征并将其投影到语言模型特征空间。

        该方法实现了完整的图像特征处理流水线，从原始像素输入到
        可用于语言模型的特征表示，是视觉-语言融合的关键步骤。

        参数：
            pixel_values (torch.Tensor): 输入图像像素值张量
                形状：(batch_size, channels, height, width)
                数值范围：通常为[0, 1]或[-1, 1]，取决于视觉编码器的预处理要求

        返回：
            torch.Tensor: 投影后的图像特征，可直接用于语言模型
                形状：(batch_size, tokens_per_image, hidden_size)
                其中 tokens_per_image = tokens_per_side * tokens_per_side

        处理流程：
        1. 视觉编码：使用vision_tower提取图像的视觉特征
        2. 特征投影：通过multi_modal_projector将视觉特征映射到语言模型空间
        3. 返回对齐后的特征，供后续的跨模态融合使用

        技术细节：
        - 视觉编码器输出密集的图像块特征序列
        - 投影器实现空间下采样和维度变换
        - 输出特征与语言模型的token表示在相同语义空间
        """
        # 步骤1：视觉编码 - 提取图像的视觉特征
        # 输入：(batch_size, channels, height, width)
        # 输出：(batch_size, num_patches, vision_hidden_size)
        vision_outputs = self.vision_tower(pixel_values)

        # 步骤2：特征投影 - 将视觉特征映射到语言模型空间
        # 输入：(batch_size, num_patches, vision_hidden_size)
        # 输出：(batch_size, tokens_per_image, hidden_size)
        image_features = self.multi_modal_projector(vision_outputs)

        return image_features

    def get_input_embeddings(
            self,
            input_ids: torch.Tensor,
            attention_mask: Optional[torch.Tensor] = None,
            **kwargs,
    ):
        """
        获取模型的输入嵌入表示，支持图像和文本的混合输入。

        该方法重写了父类的嵌入获取方法，增加了图像特征的处理能力。
        它能够识别输入序列中的图像特殊token，并将其替换为对应的图像特征，
        实现文本和图像在输入层面的统一表示。

        参数：
            input_ids (torch.Tensor): 输入token ID序列
                形状：(batch_size, seq_len)
                包含文本token和图像特殊token（vlm_config.image_tok）
            attention_mask (Optional[torch.Tensor]): 注意力掩码
                形状：(batch_size, seq_len)
                用于指示哪些位置需要参与注意力计算
            **kwargs: 其他参数，主要包括：
                - pixel_values: 图像像素值张量，形状为(batch_size, channels, height, width)

        返回：
            torch.Tensor: 混合模态的输入嵌入表示
                形状：(batch_size, seq_len, hidden_size)
                文本位置对应词嵌入，图像位置对应图像特征

        处理逻辑：
        1. 文本嵌入：将input_ids转换为词嵌入向量
        2. 图像检测：检查是否存在图像输入（pixel_values）
        3. 特征提取：如果存在图像，提取并投影图像特征
        4. 位置定位：识别输入序列中的图像特殊token位置
        5. 特征替换：将图像特殊token对应的嵌入替换为图像特征
        6. 返回混合嵌入，供语言模型处理

        技术特点：
        - 支持灵活的图文混合输入格式
        - 使用特殊token标识图像位置，便于序列处理
        - 保持与纯文本模型的接口兼容性
        - 自动处理设备和数据类型对齐
        """
        # 注释掉的代码：处理图像token超出词汇表范围的情况
        # if input_ids is not None and self.vlm_config.image_tok >= self.config.vocab_size:
        #     special_image_mask = input_ids == self.vlm_config.image_tok
        #     llm_input_ids = input_ids.clone()
        #     llm_input_ids[special_image_mask] = self.vlm_config.padding_tok
        # else:
        #     llm_input_ids = input_ids

        # 步骤1：文本嵌入 - 将token ID转换为词嵌入向量
        # 输入：(batch_size, seq_len)
        # 输出：(batch_size, seq_len, hidden_size)
        inputs_embeds = self.embed_tokens(input_ids)

        # 步骤2：图像检测 - 检查是否存在图像输入
        # pixel_values形状：(batch_size, channels, height, width)
        pixel_values: Optional[torch.Tensor] = kwargs.get('pixel_values', None)

        if pixel_values is not None:
            # 步骤3：特征提取 - 提取并投影图像特征
            # 输出：(batch_size, tokens_per_side*tokens_per_side, hidden_size)
            image_features = self.get_image_features(pixel_values)

            # 步骤4：位置定位 - 识别输入序列中的图像特殊token位置
            # 创建图像token的掩码：(batch_size, seq_len, 1)
            special_image_mask = (input_ids == self.vlm_config.image_tok).unsqueeze(-1)
            # 扩展掩码维度以匹配嵌入张量：(batch_size, seq_len, hidden_size)
            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(inputs_embeds.device)

            # 步骤5：设备对齐 - 确保图像特征与文本嵌入在同一设备和数据类型上
            image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)

            # 步骤6：特征替换 - 将图像特殊token对应的嵌入替换为图像特征
            # 使用masked_scatter实现条件替换，保持其他位置的文本嵌入不变
            # 输出：(batch_size, seq_len, hidden_size)
            inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)

        return inputs_embeds